#include "PathController.h"

MilestonePathController::MilestonePathController(Robot& robot)
  :MilestoneTrackingController(robot),path(robot)
{}

void MilestonePathController::SetSpeedLimits(Real _velScale,Real _accScale)
{
  velScale = _velScale;
  accScale = _accScale;
  path.velMax.mul(robot.velMax,velScale);
  path.accMax.mul(robot.accMax,accScale);
}

void MilestonePathController::SetPath(const vector<Vector>& _path)
{
  Assert(!_path.empty());
  path.SetMilestones(_path);
  if(xcur == path.ramps[0].x0 && dxcur == path.ramps[0].dx0) {
    MilestoneTrackingController::ramp = path.ramps[0];
    MilestoneTrackingController::pathParameter=0;
    path.ramps.erase(path.ramps.begin());
  }
  else SetMilestone(path.ramps[0].x0,path.ramps[0].dx0);
}


void MilestonePathController::SetPath(const vector<Vector>& _path,const vector<Vector>& _pathDeriv)
{
  Assert(!_path.empty());
  Assert(_path.size() == _pathDeriv.size());
  path.SetMilestones(_path,_pathDeriv);
  if(xcur == path.ramps[0].x0 && dxcur == path.ramps[0].dx0) {
    MilestoneTrackingController::ramp = path.ramps[0];
    MilestoneTrackingController::pathParameter=0;
    path.ramps.erase(path.ramps.begin());
  }
  else SetMilestone(path.ramps[0].x0,path.ramps[0].dx0);
}

void MilestonePathController::SetPath(const DynamicPath& _path)
{
  path = _path;
  if(xcur == path.ramps[0].x0 && dxcur == path.ramps[0].dx0) {
    MilestoneTrackingController::ramp = path.ramps[0];
    MilestoneTrackingController::pathParameter=0;
    path.ramps.erase(path.ramps.begin());
  }
  else SetMilestone(path.ramps[0].x0,path.ramps[0].dx0);
}

void MilestonePathController::GetPath(DynamicPath& out) const
{
  out.accMax = path.accMax;
  out.velMax = path.velMax;
  out.ramps.resize(1+path.ramps.size());
  if(MilestoneTrackingController::pathParameter < MilestoneTrackingController::ramp.endTime) {
    //normal path parameter
    out.ramps[0] = MilestoneTrackingController::ramp;
    out.ramps[0].TrimFront(MilestoneTrackingController::pathParameter);
    for(size_t i=0;i<path.ramps.size();i++)
      out.ramps[i+1]=path.ramps[i];
  }
  else { //skip the first ramp
    if(path.ramps.empty()) {
      out.ramps[0].SetConstant(MilestoneTrackingController::ramp.x1);
    }
    else {
      //trimming... this algorithm may be slow, but it shouldn't be called if all is going well
      out.ramps = path.ramps;
      Real t=MilestoneTrackingController::pathParameter-MilestoneTrackingController::ramp.endTime;
      while(!out.ramps.empty() && t >= out.ramps[0].endTime) {
	out.ramps.erase(out.ramps.begin());
	t -= out.ramps[0].endTime;
      }
      if(out.ramps.empty()) {
	out.ramps.resize(1);
	out.ramps[0].SetConstant(MilestoneTrackingController::ramp.x1);
      }
      else out.ramps[0].TrimFront(t);
      printf("GetPath() Warning: trimming off front of path: param %g, end time %g\n",MilestoneTrackingController::pathParameter,MilestoneTrackingController::ramp.endTime);
      //FatalError("GetPath() TODO: trim off front of path: param %g, end time %g\n",MilestoneTrackingController::pathParameter,MilestoneTrackingController::ramp.endTime);
    }
  }
}

const Config& MilestonePathController::Endpoint() const
{
  if(path.ramps.empty())
    return MilestoneTrackingController::Destination();
  else 
    return path.ramps.back().x1;
}

const Config& MilestonePathController::EndpointVelocity() const
{
  if(path.Empty())
    return MilestoneTrackingController::DestinationVelocity();
  else 
    return path.ramps.back().dx1;
}


void MilestonePathController::AddMilestone(const Vector& x)
{
  if(!path.Empty()) path.Append(x);
  else {
    vector<Vector> milestones(2);
    vector<Vector> dmilestones(2);
    milestones[0] = xcur;
    milestones[1] = x;
    dmilestones[0] = dxcur;
    dmilestones[1].resize(x.n,0);
    path.SetMilestones(milestones,dmilestones);
  }
}

void MilestonePathController::AddMilestone(const Vector& x,const Vector& dx)
{
  if(!path.Empty())
    path.Append(x,dx);
  else {
    vector<Vector> milestones(2);
    vector<Vector> dmilestones(2);
    milestones[0] = xcur;
    milestones[1] = x;
    dmilestones[0] = dxcur;
    dmilestones[1] = dx;
    path.SetMilestones(milestones,dmilestones);
  }
}

Real MilestonePathController::AddMilestone(const Vector& x,const Vector& dx,Real dt)
{
  size_t n=path.ramps.size();
  size_t p=n-1;
  path.ramps.resize(path.ramps.size()+1);
  if(path.ramps.size()==1) {
    Assert(dx.isZero());
    path.ramps[0].x0 = xcur;
    path.ramps[0].dx0 = dxcur;
    path.ramps[0].x1 = x;
    path.ramps[0].dx1 = dx;
  }
  else {
    path.ramps[n].x0 = path.ramps[p].x1;
    path.ramps[n].dx0 = path.ramps[p].dx1;
    path.ramps[n].x1 = x;
    path.ramps[n].dx1 = dx;
  }
  bool res=path.ramps[n].SolveMinAccel(path.velMax,dt);
  if(!res)
    res=path.ramps[n].SolveMinTime(path.accMax,path.velMax);
  Assert(res);
  return path.ramps.back().endTime;
}

void MilestonePathController::Update(Real dt)
{
  while(dt > 0) {
    if(MilestoneTrackingController::pathParameter + dt >= ramp.endTime) {
      if(path.Empty()) {
	//no path, done
	MilestoneTrackingController::Update(dt);
	return;  
      }

      //otherwise, need to move ramp forward
      Real h = MilestoneTrackingController::ramp.endTime - MilestoneTrackingController::pathParameter;
      MilestoneTrackingController::Update(h);
      dt -= h;

      //move ramp forward
      MilestoneTrackingController::ramp = path.ramps.front();
      MilestoneTrackingController::pathParameter = 0;
      path.ramps.erase(path.ramps.begin());
    }
    else {  //advancing time ends up in the same ramp, done
      MilestoneTrackingController::Update(dt);
      return; 
    }
  }
  //sanity check
  Assert(MilestoneTrackingController::pathParameter >= 0);
  Assert(path.Empty() || MilestoneTrackingController::pathParameter <= ramp.endTime);
}

void MilestonePathController::Reset()
{
  MilestoneTrackingController::Reset();
  path.ramps.clear();
}

bool MilestonePathController::ReadState(File& f)
{
  if(!MilestoneTrackingController::ReadState(f)) return false;
  int np;
  Vector temp;
  vector<Vector> milestones;
  vector<Vector> dmilestones;
  if(!ReadFile(f,np)) return false;
  for(int i=0;i<np;i++) {
    if(!ReadFile(f,temp)) return false;
    milestones.push_back(temp);
  }
  for(int i=0;i<np;i++) {
    if(!ReadFile(f,temp)) return false;
    dmilestones.push_back(temp);
  }
  path.SetMilestones(milestones,dmilestones);
  return true;
}

bool MilestonePathController::WriteState(File& f) const
{
  if(!MilestoneTrackingController::WriteState(f)) return false;
  int np=path.ramps.size();
  vector<Vector> milestones,dmilestones;
  if(np != 0) {
    np++;
    milestones.push_back(path.ramps[0].x0);
    dmilestones.push_back(path.ramps[0].dx0);
    for(size_t i=0;i<path.ramps.size();i++) {
      milestones.push_back(path.ramps[i].x1);
      dmilestones.push_back(path.ramps[i].dx1);
    }
  }
  if(!WriteFile(f,np)) return false;
  for(size_t i=0;i<milestones.size();i++)
    if(!WriteFile(f,milestones[i])) return false;
  for(size_t i=0;i<dmilestones.size();i++)
    if(!WriteFile(f,dmilestones[i])) return false;
  return true;
}
